package cn.pugle.oj.leetcode;

import cn.pugle.oj.catalog.HeapProblem;
import cn.pugle.oj.catalog.TwoProblem;

import java.util.Comparator;
import java.util.PriorityQueue;

/**
 * TODO 好题
 * https://leetcode.com/problems/find-median-from-data-stream/
 * - 插入排序
 * - 俩堆
 * - AVL树
 * - 之前分享的Quantiles
 * <p>
 * 思考, 每个数都可能成为之后的结果, 所以必须保存下来, 所以空间复杂度锁死了O(n)
 *
 * @author tzp
 * @since 2020/10/17
 */
public class LC295 implements HeapProblem, TwoProblem {

    public static class MedianFinder {

        PriorityQueue<Integer> low;
        PriorityQueue<Integer> high;

        /**
         * initialize your data structure here.
         */
        public MedianFinder() {
            low = new PriorityQueue<>(Comparator.reverseOrder());// 一定要反着用哦!!! 这里是大顶
            high = new PriorityQueue<>();// 一定要反着用哦!!! 这里是小顶
        }

        public void addNum(int num) {
            if (low.size() == 0) {
                low.add(num);
            } else {
                if (num <= low.peek()) {
                    low.add(num);
                    if (low.size() - high.size() >= 2) {
                        high.add(low.remove());
                    }
                } else {
                    high.add(num);
                    if (high.size() - low.size() >= 2) {
                        low.add(high.remove());
                    }
                }
            }
        }

        public double findMedian() {
            if (low.size() == high.size()) {
                return (low.peek() + high.peek()) / 2.0;
            } else {
                return low.size() > high.size() ? low.peek() : high.peek();
            }
        }
    }


    public static void main(String[] args) {
        MedianFinder f = new MedianFinder();
        f.addNum(1);
        f.addNum(2);
        System.out.println(f.findMedian());// -> 1.5
        f.addNum(3);
        System.out.println(f.findMedian());//-> 2
    }
}
